// Valid input consists of zero or more statements.
?start: (signature_spec | statement)*

// A signature-specification.
signature_spec: (NAME ":")? LAMBDA lambda_args "=>" signature_rhs
?signature_rhs: (signature_spec | ("[" (COLON ("," COLON)*)? "]"))


// An individual statement is either an expression or a definition.
statement: expr | definition

// A definition is equivalent to a parameter binding.
// The only difference is that definitions are top-level, i.e. not scoped.
definition: parameter_binding

// An expression is a tensorial expression or a function-definition.
?expr: tensor_expr | fn_lambda

// Function, using "lambda notation": \(param1[:], param2[:]) -> {expr}
// Clarification: The right hand side expression can indeed
// be a function-definition, indicating a lexical closure.
fn_lambda: LAMBDA lambda_args "->" expr

// lambda arguments
lambda_args: "(" (fn_arg ("," fn_arg)*)? ")"

// Tensor expression
?tensor_expr: tensor_sum | tensor_optionally_indexed

?tensor_optionally_indexed: tensor_optionally_indexed_head bracketed_indexing*

?tensor_optionally_indexed_head: tensor_paren
  | tensor_special_fn
  | tensor_atom

// A parenthesized tensor expression may evaluate to a callable,
// so can optionally have call-arguments.
// This is valid:
//   "(\(x[:], y[:]) -> x[0]*y[0])(x[:]=y[:])"
// ("We take a function of two vectors, `x` and `y`, which computes the
//  product of the leading components, and bind the 2nd input to have
//  the same value as the first, leaving us with this function:
//  "\(x[:]) -> x[0] * x[0]".
?tensor_paren: /[+-]/? "(" expr ")" call_args*

// Function calls bind zero or more parameters.
// Interpretation: if only some slots are bound, the value of the
// call-expression is the function of the still-to-be-bound named
// parameters. A parameter-binding can refer to another earlier.
// parameter-binding on the same call.
call_args: "(" (parameter_binding ("," parameter_binding)*)? ")"

// A function argument is either a bound or an unbound slot.
fn_arg: slot_spec   // Unbound slot.
  | parameter_binding    // Bound slot.

// Q: Do we handle binding function-slots against functions properly?
parameter_binding: slot_spec "=" expr

// A tensor sum-expression can start with a sign, and is a sum of one
// or more tensor product-expressions.
?tensor_sum: /[+-]/? tensor_prod (/[+-]/ tensor_prod)*

// A tensor product-expression is a parenthesized or atomic
// tensor-expression, followed by zero or more such expressions
// separated by multiplication or division.
// The operator is made mandatory, since otherwise,
// telling apart a product expression such as "x[:](y[:]+z[:])"
// from a function-application such as "x[:](y[:]=z[:])"
// may be visually confusing.
?tensor_prod: tensor_optionally_indexed (/[*\/]/ tensor_optionally_indexed)*

// An atomic tensor expression is either a number or a name,
// optionally to be called, and optionally with indexing.
?tensor_atom: number | NAME call_args*

// Index-names may get extended later to also admit transporting
// information about the group and the representation.
index_name: NAME

index_name_seq: index_name ("," index_name)*

// ==============

// "[...]" indexing where individual indices may carry tags such as "@a".
bracketed_indexing: "[" (index ("," index)*)? "]"

// An index is either a Python "start:stop:range" type index, or "d slot".
// A '+', when used as an index, corresponds to `numpy.newindex`.
// NumPy's broadcasting rules are understood, except that broadcasting
// always requires both expressions to have the same number of indices,
// and if ranges for some particular index differ, then one of them
// must be a range-1 index.
index: DIFF WS diff_slot_spec
  | stride_part? (COLON stride_part?) ~ 0..2
  | "+"


// For ease of changing the definition of some tokens.
DIFF: "d"
LAMBDA: "^"
COLON: ":"

?stride_part: int_expr


// A slot specifier, such as `pos0[:]` in ODE(pos0[:]=y[:10])
// The only allowed forms are:
// - optional list of comma-separated COLONs in square brackets.
// - literal {} to indicate that we are passing in a function here.
//   TODO: When passing in a function, do we need to do something about slot-renaming?
slot_spec: NAME /\{\}/
  | diff_slot_spec


?diff_slot_spec: NAME ("[" (COLON ("," COLON)*)? "]")?


// Integer expressions - for indexing.
// Rules parallel the tensor_{paren|sum|prod} rules.
?int_num: /[+-]?[0-9]+/

?int_atom: int_num | /[+-]/? NAME

?int_expr: int_atom | int_paren | int_sum

?int_paren: /[+-]/? "(" (int_sum | int_paren) ")"

?int_sum: /[+-]/? int_prod (/[+-]/ int_prod)*

?int_prod: (int_paren | int_atom) (/[*\/]/? (int_paren | int_atom))*


// Tensor special-functions: concatenation, stacking, reshaping,
// and other operations.
// Also: Since all "field-functions" have all slots named,
// common unnamed-slots functions (like &exp) also will be of
// this type.
// These functions all start with a '&'.
tensor_special_fn: "&" any_special_fn

?any_special_fn: special_concat
 | special_stack
 | special_reshape
 | special_onehot
 | special_zeros
 | special_es
 | special_analytic

// Tensor-concatenation special function.
// Leading parameter needs to be a literal integer,
// the axis along which to concatenate.
special_concat: "concat" "(" int_num "," tensor_expr ("," tensor_expr)* ")"

// Stacking. Syntactically equivalent to concatenation.
special_stack: "stack" "(" int_num "," tensor_expr ("," tensor_expr)* ")"

// Reshaping. Indices are to be provided as a list of integers.
special_reshape: "reshape" "(" tensor_expr "," int_num ("," int_num)* ")"

// One-hot. Index and dimension can be int-expressions.
special_onehot: "onehot" "(" int_expr "," int_expr ")"

// Vector of zeros.
special_zeros: "zeros" "(" (NAME | int_num) ")"

// Generalized "Einstein Summation". This uses rather special syntax.
// - Product-factors are separated with semicolons,
// - Multi-index product-factors are expected to have all indices named
//   with "@ indexname1, indexname2, ...".
// - 0-index (scalar) tensors need not have index names.
// - result-indices are listed after '->'.
// Interpretation:
// - multi-index array indices are named according to the order in which
//   indices occur.
// - Every right hand side index combination addresses one coefficient
//   of the resulting tensor (in index order). The left hand side
//   indices that do not also occur on the right hand side
//   are being summed over. (This "generalized" Einstein summation is
//   group-theoretically unsound, but useful for formulas that are
//   ultimately not using group theory, but short-cuts.)
special_es: "es" "(" index_named ( ";" index_named ) * "->" index_names? ")"

// Single-argument special-functions,
special_analytic: special_analytic_1arg | special_analytic_2arg | special_matrix

// For mult-index tensorial quantities, these are understood to operate
// element-wise on the coordinates.
?special_analytic_1arg: ("exp" | "log") "(" tensor_expr ")"
?special_analytic_2arg: "pow" "(" tensor_expr "," tensor_expr ")"

// Eigenvalues/Eigenvectors, 2-index-to-1-index or 2-index-to-2-index:
?special_matrix: ("eigvals" | "eigvecs") "(" tensor_expr ")"


index_named: tensor_expr ("@" index_names?)?

index_names: index_name ("," index_name)*

NAME: /[A-Za-z][A-Za-z0-9_]*/

// TODO: This needs further refinement.
number: /[0-9]+/

// Whitespace -- will get ignored.
WS: /(\s|#.*)+/

%ignore WS
